package com.qf.ability.gateway.utils;

import com.ken.event.commons.utils.ApplicationUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;

import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.TimeUnit;

/**
 * 令牌桶的工具类 - Redis key - Hash结构
 */
@Slf4j
public class TokenBucket {

    /**
     * 令牌桶的key - 通过不同的key 实现不同的限流
     */
    private String key;

    /**
     * 令牌桶的剩余令牌数
     */
    private Integer tokens;

    /**
     * 最大令牌数
     */
    private Integer maxTokens;

    /**
     * 每秒生成多少令牌
     */
    private Integer tokensSec;

    //redis的操作
    private StringRedisTemplate stringRedisTemplate;


    /**
     * 获取令牌的lua脚本
     */
    private String getTokenLua = "--获取对应的令牌桶的key\n" +
            "local key = 'token-bucket-'..KEYS[1]\n" +
            "--获取当前需要从令牌桶中拿到几个令牌\n" +
            "--tonumber将字符串转换成数字\n" +
            "local getTokens = tonumber(ARGV[1])\n" +
            "\n" +
            "--当前令牌桶中的相关参数\n" +
            "--获取当前令牌桶的剩余令牌数\n" +
            "local tokens = tonumber(redis.call('hget', key, 'tokens')) \n" +
            "--获取最大令牌数\n" +
            "local maxTokens = tonumber(redis.call('hget', key, 'maxTokens'))\n" +
            "--获取每秒生成的令牌数\n" +
            "local tokensSec = tonumber(redis.call('hget', key, 'tokensSec'))\n" +
            "--获取令牌的生成时间（微秒）\n" +
            "local tokensTime = tonumber(redis.call('hget', key, 'tokensTime'))\n" +
            "\n" +
            "--获取当前的系统时间\n" +
            "local nowTimes = redis.call('time')\n" +
            "local nowMics = tonumber(nowTimes[1]) * 1000000 + tonumber(nowTimes[2])\n" +
            "\n" +
            "--基于系统时间 重新计算令牌数\n" +
            "if nowMics > tokensTime then\n" +
            "   --只有当前时间 比 上次计算令牌的时间 要大  才会有新的令牌产生\n" +
            "   --计算生成一个令牌要多长时间\n" +
            "   local oneTokensTime = 1000000 / tokensSec\n" +
            "   --当前时间差之内，可以生成多少令牌\n" +
            "   local createTokens = (nowMics - tokensTime) / oneTokensTime\n" +
            "\n" +
            "   --计算最新剩余的令牌数\n" +
            "   tokens = math.min(tokens + createTokens, maxTokens) \t\n" +
            "   --更新令牌生成的时间\n" +
            "   tokensTime = nowMics\n" +
            "end\n" +
            "\n" +
            "--获取令牌\n" +
            "if tokens < getTokens then\n" +
            "   --令牌不足，返回-1\n" +
            "   return -1\n" +
            "end\n" +
            "\n" +
            "--说明令牌足够\n" +
            "tokens = tokens - getTokens\n" +
            "--更新令牌数\n" +
            "redis.call('hmset', key, 'tokens', tokens, 'tokensTime', tokensTime)\n" +
            "--返回成功\n" +
            "return 1";

    /**
     *  初始化令牌桶
     */
    public TokenBucket(String key, Integer maxTokens, Integer tokensSec) {
        this.key = key;
        this.tokens = maxTokens;
        this.maxTokens = maxTokens;
        this.tokensSec = tokensSec;
        this.stringRedisTemplate = ApplicationUtils.getBean(StringRedisTemplate.class);

        //初始化redis
        initRedis();
    }

    /**
     * 初始化Redis中的hash结构
     */
    public void initRedis(){

        Map<String, String> params = new HashMap<>();
        params.put("tokens", this.tokens + ""); //令牌桶的剩余令牌数
        params.put("maxTokens", this.maxTokens + ""); //令牌桶的最大令牌数
        params.put("tokensSec", this.tokensSec + ""); //每秒生成的令牌数
        params.put("tokensTime", TimeUnit.MILLISECONDS.toMicros(System.currentTimeMillis()) + "");//令牌数生成的对应时间（微秒值）

        //设置Redis的令牌桶
        stringRedisTemplate.opsForHash().putAll("token-bucket-" + this.key, params);

        log.debug("[token bucket] 令牌桶初始化完成【{} - {}】", this.key, params);
    }


    /**
     * 获取令牌
     *  getTokens - 需要获取的令牌
     * @return 如果返回true - 表示已经获取到对应令牌，可以放行  如果返回false - 表示没有获得令牌 直接拒绝
     */
    public boolean getTokens(Integer getTokens){
        //通过lua脚本 实现令牌的获取 - 因为redis单线程的原因 可以保证获取令牌时 是线程安全的
        Long result = stringRedisTemplate.execute(new DefaultRedisScript<>(getTokenLua, Long.class),
                Collections.singletonList(this.key), // lua -> KEYS[]
                getTokens + "");// lua -> ARGV[]

        return result == 1;
    }
}